package info.wangyuan.agent.service.curd.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.IoUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import info.wangyuan.agent.common.Constant.FileConstant;
import info.wangyuan.agent.common.Enum.FileEventType;
import info.wangyuan.agent.entity.po.KnowledgeBase;
import info.wangyuan.agent.exception.BusinessException;
import info.wangyuan.agent.mapper.KnowledgeBaseMapper;
import info.wangyuan.agent.service.curd.CommonService;
import info.wangyuan.agent.service.curd.KnowledgeBaseService;
import info.wangyuan.agent.service.event.RagFileEvent;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
 * 知识库服务实现类
 *
 * @author Albert
 * @since 2025-07-15 21:47:43
 */
@Slf4j
@Service
public class KnowledgeBaseServiceImpl extends ServiceImpl<KnowledgeBaseMapper, KnowledgeBase> implements KnowledgeBaseService {

    private final CommonService commonService;
    private final KnowledgeBaseMapper knowledgeBaseMapper;
    private final ApplicationEventPublisher publisher;


    public KnowledgeBaseServiceImpl(CommonService commonService, KnowledgeBaseMapper knowledgeBaseMapper, ApplicationEventPublisher publisher) {
        this.commonService = commonService;
        this.knowledgeBaseMapper = knowledgeBaseMapper;
        this.publisher = publisher;
    }

    /**
     * 添加知识库
     */
    @Override
    public void canAdd(KnowledgeBase kb) {
        String name = kb.getName().trim();

        if (count(new QueryWrapper<KnowledgeBase>().eq("name", name)) > 0) {
            throw new BusinessException("名称已存在");
        }
        File dir = new File(FileConstant.RAG_ROOT_PATH, name);
        if (!dir.exists()) {
            if (!dir.mkdirs()) {
                throw new BusinessException("知识库创建失败");
            }
        }
        kb.setPath(File.separator + name);
        knowledgeBaseMapper.insert(kb);
    }


    /**
     * 修改知识库
     */
    @Override
    public void canUpdate(KnowledgeBase kb) {
        String name = kb.getName().trim();
        if (count(new QueryWrapper<KnowledgeBase>()
                .ne("id", kb.getId())
                .eq("name", name)) > 0) {
            throw new BusinessException("名称已存在");
        }
        File newDir = new File(FileConstant.RAG_ROOT_PATH, name);
        KnowledgeBase dbKb = getById(kb);
        if (!newDir.exists()) {
            File oldDir = new File(FileConstant.RAG_ROOT_PATH, dbKb.getPath().substring(1));
            if (!oldDir.renameTo(newDir)) {
                throw new BusinessException("知识库重命名失败");
            }
        }
        knowledgeBaseMapper.updateById(kb);
//        vectorService.rebuildKnowledgeBaseVectorsAsync(dbKb.getName(), kb.getName(), dbKb.getPath());
    }


    /**
     * 删除知识库
     */
    @Override
    public void canDelete(Integer id) {
        getOptById(id).ifPresent(kb -> {
            try {
                // 删除整个知识库文件夹
                FileUtil.del(FileConstant.RAG_ROOT_PATH + kb.getPath());
                // 从知识源中删除知识库
                commonService.deleteKnowledgeSource(id);
                // 删除知识库文件数据
                publisher.publishEvent(new RagFileEvent<>(FileEventType.REMOVE_BY_DATASET_ID, id));
            } catch (Exception e) {
                log.error("知识库删除失败，原因：{}", e.getMessage(), e);
                throw new BusinessException("知识库删除失败");
            }
        });
    }


    /**
     * 压缩并下载知识库
     */
    @Override
    public ResponseEntity<StreamingResponseBody> downloadFolderZip(Integer id) throws IOException {
        KnowledgeBase kb = getById(id);
        Path folderPath = Paths.get(FileConstant.RAG_ROOT_PATH + kb.getPath());

        // StreamingResponseBody用来流式传输
        StreamingResponseBody stream = out -> {
            try (ZipOutputStream zos = new ZipOutputStream(out)) {
                Files.walk(folderPath)
                        .filter(path -> !Files.isDirectory(path))
                        .forEach(path -> {
                            try (InputStream is = Files.newInputStream(path)) {
                                zos.putNextEntry(new ZipEntry(folderPath.relativize(path).toString()));
                                IoUtil.copy(is, zos);
                                zos.closeEntry();
                            } catch (IOException e) {
                                log.error("压缩文件失败：{}", e.getMessage());
                                throw new BusinessException("压缩失败，无法处理文件：" + path.toString());
                            }
                        });
            } catch (IOException e) {
                log.error("ZIP 文件生成失败：{}", e.getMessage());
                throw new BusinessException("压缩失败，无法创建ZIP文件");
            }
        };

        // 获取文件夹总大小
        long folderSize = Files.walk(folderPath)
                .filter(path -> !Files.isDirectory(path))
                .mapToLong(path -> path.toFile().length())
                .sum();

        // 设置响应头并返回流式文件
        return ResponseEntity.ok()
                .header(
                        HttpHeaders.CONTENT_DISPOSITION,
                        "attachment; filename=" + URLEncoder.encode(kb.getName(), StandardCharsets.UTF_8) + ".zip"
                )
                .contentLength(folderSize)
                .contentType(MediaType.APPLICATION_OCTET_STREAM)
                .body(stream);
    }

    /**
     * 批量校验知识库是否存在
     */
    @Override
    public void checkKnowledgeBasesExist(List<Integer> knowledgeBaseIds) {
        if (CollUtil.isEmpty(knowledgeBaseIds)) {
            throw new IllegalArgumentException("必须指定知识源");
        }

        // 一次性查数据库
        List<Integer> existIds = knowledgeBaseMapper.selectList(
                Wrappers.lambdaQuery(KnowledgeBase.class)
                        .in(KnowledgeBase::getId, knowledgeBaseIds)
                        .select(KnowledgeBase::getId)
        ).stream().map(KnowledgeBase::getId).toList();

        // 找出不存在的 ID
        List<Integer> notExistIds = knowledgeBaseIds.stream()
                .filter(id -> !existIds.contains(id))
                .toList();

        if (!notExistIds.isEmpty()) {
            throw new BusinessException("以下知识库不存在: " + notExistIds);
        }
    }

}
